
# ---------------------------------------------
# Copyright (c) OpenMMLab. All rights reserved.
# ---------------------------------------------
#  Modified by Zhiqi Li
# ---------------------------------------------

from mmcv.ops.multi_scale_deform_attn import multi_scale_deformable_attn_pytorch
import warnings
import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import xavier_init, constant_init
from mmcv.cnn.bricks.registry import (ATTENTION,
                                      TRANSFORMER_LAYER,
                                      TRANSFORMER_LAYER_SEQUENCE)
from mmcv.cnn.bricks.transformer import build_attention
import math
from mmcv.runner import force_fp32, auto_fp16

from mmcv.runner.base_module import BaseModule, ModuleList, Sequential

from mmcv.utils import ext_loader
from .multi_scale_deformable_attn_function import MultiScaleDeformableAttnFunction_fp32, \
    MultiScaleDeformableAttnFunction_fp16
from projects.mmdet3d_plugin.models.utils.bricks import run_time
ext_module = ext_loader.load_ext(
    '_ext', ['ms_deform_attn_backward', 'ms_deform_attn_forward'])


@ATTENTION.register_module()
class SpatialCrossAttention(BaseModule):
    """An attention module used in BEVFormer.
    Args:
        embed_dims (int): The embedding dimension of Attention.
            Default: 256.
        num_cams (int): The number of cameras
        dropout (float): A Dropout layer on `inp_residual`.
            Default: 0..
        init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
            Default: None.
        deformable_attention: (dict): The config for the deformable attention used in SCA.
    """

    def __init__(self,
                 embed_dims=256,
                 num_cams=6,
                 pc_range=None,
                 dropout=0.1,
                 init_cfg=None,
                 batch_first=False,
                 deformable_attention=dict(
                     type='MSDeformableAttention3D',
                     embed_dims=256,
                     num_levels=4),
                 **kwargs
                 ):
        super(SpatialCrossAttention, self).__init__(init_cfg)

        self.init_cfg = init_cfg
        self.dropout = nn.Dropout(dropout)
        self.pc_range = pc_range
        self.fp16_enabled = False
        self.deformable_attention = build_attention(deformable_attention)
        self.embed_dims = embed_dims
        self.num_cams = num_cams
        self.output_proj = nn.Linear(embed_dims, embed_dims)
        self.batch_first = batch_first
        self.init_weight()

    def init_weight(self):
        """Default initialization for Parameters of Module."""
        xavier_init(self.output_proj, distribution='uniform', bias=0.)
    
    @force_fp32(apply_to=('query', 'key', 'value', 'query_pos', 'reference_points_cam'))
    def forward(self,
                query,
                key,
                value,
                residual=None,
                query_pos=None,
                key_padding_mask=None,
                reference_points=None,
                spatial_shapes=None,
                reference_points_cam=None,
                bev_mask=None,
                level_start_index=None,
                flag='encoder',
                **kwargs):
        """
        前向传播函数，实现 BEVFormer 中的空间交叉注意力机制（Spatial Cross Attention）。
        该模块的主要功能是：对于每个 BEV 查询（BEV query），仅与其对应的摄像头视角下的特定 BEV 区域进行交互，从而节省显存并提升效率。

        变量说明及维度：
        query: (bs, num_query, embed_dims) 查询特征
        key: (num_cams, l, bs, embed_dims) 键特征
        value: (num_cams, l, bs, embed_dims) 值特征
        residual: (bs, num_query, embed_dims) 残差连接
        query_pos: (bs, num_query, embed_dims) 查询位置编码
        key_padding_mask: 未使用
        reference_points: (bs, num_query, 4) 归一化参考点
        spatial_shapes: (num_levels, 2) 多尺度特征的空间尺寸
        reference_points_cam: (num_cams, bs, num_query, D, 2) 每个摄像头下的参考点
        bev_mask: (num_cams, bs, num_query, H*W) BEV 区域掩码
        level_start_index: (num_levels,) 每个尺度起始索引
        flag: 字符串，未使用
        kwargs: 其他参数

        返回：
        (bs, num_query, embed_dims) 输出特征
        """

        if key is None:
            key = query
        if value is None:
            value = key

        if residual is None:
            inp_residual = query
            slots = torch.zeros_like(query)
        if query_pos is not None:
            query = query + query_pos

        bs, num_query, _ = query.size()

        D = reference_points_cam.size(3)  # D: 参考点在深度方向的数量
        indexes = []
        for i, mask_per_img in enumerate(bev_mask):
            # mask_per_img: (bs, num_query, H*W)
            index_query_per_img = mask_per_img[0].sum(-1).nonzero().squeeze(-1)
            # index_query_per_img: (有效 query 的索引)
            indexes.append(index_query_per_img)
        max_len = max([len(each) for each in indexes])

        # 每个摄像头只与其对应的 BEV 查询交互，节省显存
        queries_rebatch = query.new_zeros(
            [bs, self.num_cams, max_len, self.embed_dims])  # (bs, num_cams, max_len, embed_dims)
        reference_points_rebatch = reference_points_cam.new_zeros(
            [bs, self.num_cams, max_len, D, 2])  # (bs, num_cams, max_len, D, 2)
        
        for j in range(bs):
            for i, reference_points_per_img in enumerate(reference_points_cam):   
                index_query_per_img = indexes[i]
                queries_rebatch[j, i, :len(index_query_per_img)] = query[j, index_query_per_img]
                reference_points_rebatch[j, i, :len(index_query_per_img)] = reference_points_per_img[j, index_query_per_img]

        num_cams, l, bs, embed_dims = key.shape

        # 变换 key 和 value 的维度以适配 deformable attention
        key = key.permute(2, 0, 1, 3).reshape(
            bs * self.num_cams, l, self.embed_dims)  # (bs*num_cams, l, embed_dims)
        value = value.permute(2, 0, 1, 3).reshape(
            bs * self.num_cams, l, self.embed_dims)  # (bs*num_cams, l, embed_dims)

        # 调用可变形注意力模块
        queries = self.deformable_attention(
            query=queries_rebatch.view(bs*self.num_cams, max_len, self.embed_dims),
            key=key,
            value=value,
            reference_points=reference_points_rebatch.view(bs*self.num_cams, max_len, D, 2),
            spatial_shapes=spatial_shapes,
            level_start_index=level_start_index
        ).view(bs, self.num_cams, max_len, self.embed_dims)  # (bs, num_cams, max_len, embed_dims)

        # 将每个摄像头的结果累加到对应的 BEV 查询上
        for j in range(bs):
            for i, index_query_per_img in enumerate(indexes):
                slots[j, index_query_per_img] += queries[j, i, :len(index_query_per_img)]

        # 统计每个 BEV 查询被多少个摄像头覆盖
        count = bev_mask.sum(-1) > 0  # (num_cams, bs, num_query)
        count = count.permute(1, 2, 0).sum(-1)  # (bs, num_query)
        count = torch.clamp(count, min=1.0)
        slots = slots / count[..., None]  # 归一化
        slots = self.output_proj(slots)   # 线性投影

        return self.dropout(slots) + inp_residual  # 残差连接


@ATTENTION.register_module()
class MSDeformableAttention3D(BaseModule):
    """An attention module used in BEVFormer based on Deformable-Detr.
    `Deformable DETR: Deformable Transformers for End-to-End Object Detection.
    <https://arxiv.org/pdf/2010.04159.pdf>`_.
    Args:
        embed_dims (int): The embedding dimension of Attention.
            Default: 256.
        num_heads (int): Parallel attention heads. Default: 64.
        num_levels (int): The number of feature map used in
            Attention. Default: 4.
        num_points (int): The number of sampling points for
            each query in each head. Default: 4.
        im2col_step (int): The step used in image_to_column.
            Default: 64.
        dropout (float): A Dropout layer on `inp_identity`.
            Default: 0.1.
        batch_first (bool): Key, Query and Value are shape of
            (batch, n, embed_dim)
            or (n, batch, embed_dim). Default to False.
        norm_cfg (dict): Config dict for normalization layer.
            Default: None.
        init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
            Default: None.
    """

    def __init__(self,
                 embed_dims=256,
                 num_heads=8,
                 num_levels=4,
                 num_points=8,
                 im2col_step=64,
                 dropout=0.1,
                 batch_first=True,
                 norm_cfg=None,
                 init_cfg=None):
        super().__init__(init_cfg)
        if embed_dims % num_heads != 0:
            raise ValueError(f'embed_dims must be divisible by num_heads, '
                             f'but got {embed_dims} and {num_heads}')
        dim_per_head = embed_dims // num_heads
        self.norm_cfg = norm_cfg
        self.batch_first = batch_first
        self.output_proj = None
        self.fp16_enabled = False

        # you'd better set dim_per_head to a power of 2
        # which is more efficient in the CUDA implementation
        def _is_power_of_2(n):
            if (not isinstance(n, int)) or (n < 0):
                raise ValueError(
                    'invalid input for _is_power_of_2: {} (type: {})'.format(
                        n, type(n)))
            return (n & (n - 1) == 0) and n != 0

        if not _is_power_of_2(dim_per_head):
            warnings.warn(
                "You'd better set embed_dims in "
                'MultiScaleDeformAttention to make '
                'the dimension of each attention head a power of 2 '
                'which is more efficient in our CUDA implementation.')

        self.im2col_step = im2col_step
        self.embed_dims = embed_dims
        self.num_levels = num_levels
        self.num_heads = num_heads
        self.num_points = num_points
        self.sampling_offsets = nn.Linear(
            embed_dims, num_heads * num_levels * num_points * 2)
        self.attention_weights = nn.Linear(embed_dims,
                                           num_heads * num_levels * num_points)
        self.value_proj = nn.Linear(embed_dims, embed_dims)

        self.init_weights()

    def init_weights(self):
        """Default initialization for Parameters of Module."""
        constant_init(self.sampling_offsets, 0.)
        thetas = torch.arange(
            self.num_heads,
            dtype=torch.float32) * (2.0 * math.pi / self.num_heads)
        grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
        grid_init = (grid_init /
                     grid_init.abs().max(-1, keepdim=True)[0]).view(
            self.num_heads, 1, 1,
            2).repeat(1, self.num_levels, self.num_points, 1)
        for i in range(self.num_points):
            grid_init[:, :, i, :] *= i + 1

        self.sampling_offsets.bias.data = grid_init.view(-1)
        constant_init(self.attention_weights, val=0., bias=0.)
        xavier_init(self.value_proj, distribution='uniform', bias=0.)
        xavier_init(self.output_proj, distribution='uniform', bias=0.)
        self._is_init = True

    def forward(self,
                query,
                key=None,
                value=None,
                identity=None,
                query_pos=None,
                key_padding_mask=None,
                reference_points=None,
                spatial_shapes=None,
                level_start_index=None,
                **kwargs):
        """
        MultiScaleDeformAttention 的前向传播函数。
        实现了多尺度可变形注意力机制，用于对输入的 query 进行特征增强。

        参数说明：
        query (Tensor): 查询特征，形状为 (bs, num_query, embed_dims)。
        key (Tensor): 键特征，形状为 (bs, num_key, embed_dims)，未用到。
        value (Tensor): 值特征，形状为 (bs, num_key, embed_dims)。
        identity (Tensor): 残差连接，形状与 query 相同，默认用 query。
        query_pos (Tensor): 查询位置编码。
        key_padding_mask (Tensor): 键的掩码，形状为 (bs, num_key)。
        reference_points (Tensor): 归一化参考点，形状为 (bs, num_query, num_levels, 2) 或 (bs, num_query, num_levels, 4)。
        spatial_shapes (Tensor): 每个特征层的空间尺寸，形状为 (num_levels, 2)。
        level_start_index (Tensor): 每个特征层的起始索引，形状为 (num_levels,)。
        kwargs: 其他参数。

        返回：
        Tensor: 输出特征，形状为 (bs, num_query, embed_dims)。
        """

        # 如果未提供 value，则用 query
        if value is None:
            value = query
        # 如果未提供 identity，则用 query
        if identity is None:
            identity = query
        # 如果有位置编码，则加到 query 上
        if query_pos is not None:
            query = query + query_pos

        # 如果不是 batch_first 格式，则转为 (bs, num_query, embed_dims)
        if not self.batch_first:
            query = query.permute(1, 0, 2)
            value = value.permute(1, 0, 2)

        bs, num_query, _ = query.shape
        bs, num_value, _ = value.shape
        # 检查 value 的长度是否等于所有特征层空间尺寸之积
        assert (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == num_value

        # 对 value 做线性变换
        value = self.value_proj(value)
        # 如果有掩码，则将被 mask 的位置置零
        if key_padding_mask is not None:
            value = value.masked_fill(key_padding_mask[..., None], 0.0)
        # (bs, num_value, num_heads, head_dim)
        value = value.view(bs, num_value, self.num_heads, -1)
        # 计算采样偏移量 (bs, num_query, num_heads, num_levels, num_points, 2)
        sampling_offsets = self.sampling_offsets(query).view(
            bs, num_query, self.num_heads, self.num_levels, self.num_points, 2)
        # 计算注意力权重 (bs, num_query, num_heads, num_levels * num_points)
        attention_weights = self.attention_weights(query).view(
            bs, num_query, self.num_heads, self.num_levels * self.num_points)
        # softmax 归一化
        attention_weights = attention_weights.softmax(-1)
        # (bs, num_query, num_heads, num_levels, num_points)
        attention_weights = attention_weights.view(bs, num_query,
                                                  self.num_heads,
                                                  self.num_levels,
                                                  self.num_points)

        # 处理参考点
        if reference_points.shape[-1] == 2:
            """
            对于每个 BEV 查询，在 3D 空间有 num_Z_anchors 个锚点（不同高度）。
            投影后，每个 BEV 查询在每个 2D 图像有 num_Z_anchors 个参考点。
            每个参考点采样 num_points 个采样点。
            总采样点数为 num_points * num_Z_anchors。
            """
            # 归一化偏移量
            offset_normalizer = torch.stack(
                [spatial_shapes[..., 1], spatial_shapes[..., 0]], -1)

            bs, num_query, num_Z_anchors, xy = reference_points.shape
            # 扩展维度以便后续相加
            reference_points = reference_points[:, :, None, None, None, :, :]
            sampling_offsets = sampling_offsets / \
                offset_normalizer[None, None, None, :, None, :]
            bs, num_query, num_heads, num_levels, num_all_points, xy = sampling_offsets.shape
            # 重新排列采样偏移量
            sampling_offsets = sampling_offsets.view(
                bs, num_query, num_heads, num_levels, num_all_points // num_Z_anchors, num_Z_anchors, xy)
            # 得到采样位置
            sampling_locations = reference_points + sampling_offsets
            bs, num_query, num_heads, num_levels, num_points, num_Z_anchors, xy = sampling_locations.shape
            assert num_all_points == num_points * num_Z_anchors

            # 展平成 (bs, num_query, num_heads, num_levels, num_all_points, 2)
            sampling_locations = sampling_locations.view(
                bs, num_query, num_heads, num_levels, num_all_points, xy)

        elif reference_points.shape[-1] == 4:
            # 暂不支持 4 维参考点
            assert False
        else:
            raise ValueError(
                f'Last dim of reference_points must be'
                f' 2 or 4, but get {reference_points.shape[-1]} instead.')

        # 采样位置: (bs, num_query, num_heads, num_levels, num_all_points, 2)
        # 注意力权重: (bs, num_query, num_heads, num_levels, num_all_points)

        # 调用 CUDA 实现或 PyTorch 实现的多尺度可变形注意力
        if torch.cuda.is_available() and value.is_cuda:
            if value.dtype == torch.float16:
                MultiScaleDeformableAttnFunction = MultiScaleDeformableAttnFunction_fp32
            else:
                MultiScaleDeformableAttnFunction = MultiScaleDeformableAttnFunction_fp32
            output = MultiScaleDeformableAttnFunction.apply(
                value, spatial_shapes, level_start_index, sampling_locations,
                attention_weights, self.im2col_step)
        else:
            output = multi_scale_deformable_attn_pytorch(
                value, spatial_shapes, sampling_locations, attention_weights)
        # 如果不是 batch_first，则转回 (num_query, bs, embed_dims)
        if not self.batch_first:
            output = output.permute(1, 0, 2)

        return output
